import os
import re
import sys
from urllib.parse import urlparse, urlsplit

from download_utils import download_file_with_hf
from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor, as_completed

# 设置模型后缀
MODEL_EXTENSIONS = ('.safetensors', '.ckpt', '.pt', '.pth')

def is_model_url(url):
    path = urlsplit(url).path.lower()
    return any(path.endswith(ext) for ext in MODEL_EXTENSIONS)

def extract_links_with_paths(text):
    lines = text.splitlines()
    path_pattern = r"- `(/[^`]+)`"
    url_pattern = r"\[([^\]]+)\]\((https?://[^\)]+)\)"

    result = defaultdict(list)
    current_path = None

    for line in lines:
        path_match = re.match(path_pattern, line.strip())
        if path_match:
            current_path = path_match.group(1)
        elif current_path:
            url_match = re.findall(url_pattern, line.strip())
            for filename, url in url_match:
                if is_model_url(url):
                    result[current_path].append((filename, url))
                    # print(f"current_path{current_path}==={filename}=={url}")

    return result

def download(url, filename, save_path):
    try:
        print(f"使用 huggingface_hub 下载: {url}")
        download_file_with_hf(url, filename, save_path)
        print(f"✅ 下载完成: {filename}")
    except Exception as e:
        print(f"❌ 下载失败: {filename}，错误: {e}")

def main():
    if len(sys.argv) != 2:
        print("用法: python down.py <README.md 文件路径>")
        sys.exit(1)

    input_file = sys.argv[1]
    save_root = "./downmodels"

    print(f"📂 正在处理文件: {input_file}")
    with open(input_file, "r", encoding="utf-8") as f:
        input_text = f.read()

    entries = extract_links_with_paths(input_text)

    if not entries:
        print("⚠ 未找到任何模型链接")
        return

    with ThreadPoolExecutor(max_workers=5) as executor:
        futures = []
        for path, files in entries.items():
            print(f"📁 路径: {path}")
            save_path = os.path.join(save_root, path.lstrip('/'))
            for filename, url in files:
                print(f"  📄 文件: {filename} | URL: {url}")
                futures.append(executor.submit(download, url, filename, save_path))
        
        for future in as_completed(futures):
            try:
                future.result()
            except Exception as e:
                print(f"下载任务失败: {e}")

if __name__ == "__main__":
    main()
